"""
Hybrid PETSc-Kokkos Builder Implementation.

This module provides the HybridBuilder class that wraps the build_hybrid.sh 
script and integrates it with the GSM extension system.
"""

import os
import shutil
import subprocess
import time
import logging
from pathlib import Path
from typing import Dict, List, Optional, Any

from gcr_solver_manager.extensions.base_builder import (
    BuilderInterface, BuildConfig, BuildResult, BuildStatus, BuildDependencyResult
)


class HybridBuilder(BuilderInterface):
    """Builder for Hybrid PETSc-Kokkos framework."""
    
    def __init__(self):
        """Initialize the hybrid builder."""
        self.logger = logging.getLogger(__name__)
        self._hybrid_root = Path(__file__).parent.parent.parent.parent.parent / "src" / "prototypes" / "hybrid"
        self._build_script = self._hybrid_root / "build_hybrid.sh"
        
    @property
    def hybrid_root(self) -> Path:
        """Get the hybrid framework root directory."""
        return self._hybrid_root
        
    @property 
    def build_script(self) -> Path:
        """Get the build script path."""
        return self._build_script
    
    def check_dependencies(self) -> BuildDependencyResult:
        """Check if required dependencies are available."""
        self.logger.info("Checking hybrid framework dependencies...")
        
        missing_deps = []
        version_conflicts = []
        warnings = []
        dependency_info = {}
        
        # Check build script exists
        if not self._build_script.exists():
            missing_deps.append(f"Build script not found: {self._build_script}")
            return BuildDependencyResult(
                all_satisfied=False,
                missing_dependencies=missing_deps,
                version_conflicts=version_conflicts,
                warnings=warnings,
                dependency_info=dependency_info
            )
        
        # Check CMake
        try:
            result = subprocess.run(["cmake", "--version"], capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                version_line = result.stdout.split('\n')[0]
                cmake_version = version_line.split()[2]
                dependency_info["cmake"] = cmake_version
                self.logger.debug(f"Found CMake: {cmake_version}")
            else:
                missing_deps.append("cmake")
        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
            missing_deps.append("cmake")
        
        # Check MPI
        try:
            result = subprocess.run(["mpicc", "--version"], capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                mpi_info = result.stdout.split('\n')[0]
                dependency_info["mpi"] = mpi_info
                self.logger.debug(f"Found MPI: {mpi_info}")
            else:
                missing_deps.append("mpi")
        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
            missing_deps.append("mpi")
        
        # Check PETSc environment
        petsc_dir = os.environ.get("PETSC_DIR")
        if petsc_dir and Path(petsc_dir).exists():
            dependency_info["petsc_dir"] = petsc_dir
            petsc_arch = os.environ.get("PETSC_ARCH")
            if petsc_arch:
                dependency_info["petsc_arch"] = petsc_arch
            self.logger.debug(f"Found PETSc: {petsc_dir}")
        else:
            missing_deps.append("PETSC_DIR environment variable or directory")
        
        # Check Kokkos (optional, may be found by CMake)
        kokkos_root = os.environ.get("Kokkos_ROOT")
        if kokkos_root and Path(kokkos_root).exists():
            dependency_info["kokkos_root"] = kokkos_root
            self.logger.debug(f"Found Kokkos: {kokkos_root}")
        else:
            warnings.append("Kokkos_ROOT not set - CMake will search standard locations")
        
        # Check GPU backends
        gpu_backends = self._detect_gpu_backends()
        dependency_info["gpu_backends"] = gpu_backends
        if not gpu_backends:
            warnings.append("No GPU backends detected - will use CPU-only build")
        
        all_satisfied = len(missing_deps) == 0
        
        return BuildDependencyResult(
            all_satisfied=all_satisfied,
            missing_dependencies=missing_deps,
            version_conflicts=version_conflicts,
            warnings=warnings,
            dependency_info=dependency_info
        )
    
    def _detect_gpu_backends(self) -> List[str]:
        """Detect available GPU backends."""
        backends = []
        
        # Check CUDA
        try:
            result = subprocess.run(["nvcc", "--version"], capture_output=True, timeout=5)
            if result.returncode == 0:
                backends.append("cuda")
        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
            pass
        
        # Check HIP
        try:
            result = subprocess.run(["hipcc", "--version"], capture_output=True, timeout=5)
            if result.returncode == 0:
                backends.append("hip")
        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
            pass
        
        return backends
    
    def build(self, config: BuildConfig) -> BuildResult:
        """Perform framework-specific build operations."""
        self.logger.info("Starting hybrid framework build...")
        start_time = time.time()
        
        # Initialize result
        result = BuildResult(
            status=BuildStatus.FAILURE,
            start_time=start_time
        )
        
        try:
            # Validate configuration and dependencies
            dep_result = self.check_dependencies()
            if not dep_result.all_satisfied:
                result.error_message = f"Dependencies not satisfied: {dep_result.missing_dependencies}"
                return result
            
            # Build command line arguments
            build_args = self._build_command_args(config, dep_result)
            
            # Set working directory
            working_dir = config.source_dir or self._hybrid_root
            build_dir = config.build_dir or (working_dir / "build")
            
            # Setup environment
            env = os.environ.copy()
            env.update(config.environment)
            
            # Execute build script
            self.logger.info(f"Executing: {' '.join(build_args)}")
            self.logger.info(f"Working directory: {working_dir}")
            
            process = subprocess.run(
                build_args,
                cwd=working_dir,
                env=env,
                capture_output=True,
                text=True,
                timeout=3600  # 1 hour timeout
            )
            
            # Record outputs
            result.stdout = process.stdout
            result.stderr = process.stderr
            result.build_log = f"{process.stdout}\n{process.stderr}"
            result.error_code = process.returncode
            
            # Determine success
            if process.returncode == 0:
                result.status = BuildStatus.SUCCESS
                result.build_dir = build_dir
                
                # Find built binary
                binary_path = self._find_built_binary(build_dir)
                if binary_path:
                    result.binary_path = binary_path
                    self.logger.info(f"Build successful: {binary_path}")
                else:
                    result.status = BuildStatus.PARTIAL
                    result.error_message = "Build succeeded but binary not found"
                    self.logger.warning("Build completed but binary not located")
            else:
                result.status = BuildStatus.FAILURE
                result.error_message = f"Build script failed with code {process.returncode}"
                self.logger.error(f"Build failed: {result.error_message}")
                
                # Log error details
                if process.stderr:
                    self.logger.error(f"Build stderr: {process.stderr}")
        
        except subprocess.TimeoutExpired:
            result.error_message = "Build timed out after 1 hour"
            result.status = BuildStatus.FAILURE
            self.logger.error(result.error_message)
        
        except Exception as e:
            result.error_message = f"Build error: {str(e)}"
            result.status = BuildStatus.FAILURE
            self.logger.error(result.error_message, exc_info=True)
        
        finally:
            result.end_time = time.time()
            result.duration = result.end_time - start_time
        
        return result
    
    def _build_command_args(self, config: BuildConfig, dep_result: BuildDependencyResult) -> List[str]:
        """Build command line arguments for build_hybrid.sh."""
        args = [str(self._build_script)]
        
        # Build type
        if config.debug_mode:
            args.append("--build-type=Debug")
        else:
            args.append("--build-type=Release")
        
        # Clean build
        if config.clean_build:
            args.append("--clean")
        
        # Verbose output
        if config.verbose:
            args.append("--verbose")
        
        # GPU backends
        gpu_backends = dep_result.dependency_info.get("gpu_backends", [])
        if config.gpu_backend == "cuda" and "cuda" in gpu_backends:
            args.append("--enable-cuda")
        elif config.gpu_backend == "hip" and "hip" in gpu_backends:
            args.append("--enable-hip")
        elif gpu_backends and config.gpu_backend is None:
            # Auto-detect first available backend
            if "cuda" in gpu_backends:
                args.append("--enable-cuda")
            elif "hip" in gpu_backends:
                args.append("--enable-hip")
        
        # Testing and benchmarks
        if not config.enable_testing:
            args.append("--disable-testing")
        if config.enable_benchmarks:
            args.append("--enable-benchmarks")
        
        # Install prefix
        if config.install_prefix:
            args.append(f"--install-prefix={config.install_prefix}")
        
        # Additional CMake args
        for cmake_arg in config.extra_cmake_args:
            args.append(cmake_arg)
        
        return args
    
    def _find_built_binary(self, build_dir: Path) -> Optional[Path]:
        """Find the main hybrid binary in build directory."""
        possible_names = [
            "hybrid_prototype",
            "hybrid_solver",
            "hybrid_test_suite"
        ]
        
        for name in possible_names:
            binary_path = build_dir / name
            if binary_path.exists() and binary_path.is_file():
                # Check if executable
                if os.access(binary_path, os.X_OK):
                    return binary_path
        
        # Search recursively if not found at root
        for pattern in possible_names:
            matches = list(build_dir.rglob(pattern))
            for match in matches:
                if match.is_file() and os.access(match, os.X_OK):
                    return match
        
        return None
    
    def clean(self, config: BuildConfig) -> bool:
        """Clean build artifacts."""
        self.logger.info("Cleaning hybrid framework build artifacts...")
        
        try:
            build_dir = config.build_dir or (self._hybrid_root / "build")
            
            if build_dir.exists():
                shutil.rmtree(build_dir)
                self.logger.info(f"Removed build directory: {build_dir}")
            
            # Clean other common build artifacts
            artifacts_to_clean = [
                self._hybrid_root / "CMakeCache.txt",
                self._hybrid_root / "CMakeFiles",
                self._hybrid_root / "compile_commands.json"
            ]
            
            for artifact in artifacts_to_clean:
                if artifact.exists():
                    if artifact.is_file():
                        artifact.unlink()
                    else:
                        shutil.rmtree(artifact)
                    self.logger.debug(f"Removed: {artifact}")
            
            return True
            
        except Exception as e:
            self.logger.error(f"Clean failed: {e}")
            return False
    
    def get_binary_path(self, config: BuildConfig) -> Optional[Path]:
        """Get path to built binary."""
        build_dir = config.build_dir or (self._hybrid_root / "build")
        return self._find_built_binary(build_dir)
    
    def get_build_info(self, config: BuildConfig) -> Dict[str, Any]:
        """Get information about build configuration."""
        base_info = super().get_build_info(config)
        
        hybrid_info = {
            "hybrid_root": str(self._hybrid_root),
            "build_script": str(self._build_script),
            "script_exists": self._build_script.exists(),
        }
        
        # Add dependency info
        dep_result = self.check_dependencies()
        hybrid_info.update({
            "dependencies": dep_result.dependency_info,
            "dependencies_satisfied": dep_result.all_satisfied,
            "missing_dependencies": dep_result.missing_dependencies,
            "warnings": dep_result.warnings
        })
        
        base_info["hybrid_specific"] = hybrid_info
        return base_info
    
    def supports_incremental_build(self) -> bool:
        """Check if builder supports incremental builds."""
        return True  # CMake-based builds support incremental compilation